#include "hero_integrated_vision/pnp_solver.hpp"

#include <eigen3/Eigen/Dense>
#include <opencv2/calib3d.hpp>
#include <opencv2/core/eigen.hpp>
#include <vector>
#include <angles/angles.h>
#include <tf2/LinearMath/Matrix3x3.h>
#include <tf2/LinearMath/Quaternion.h>
#include <tf2/convert.h>
#include <tf2_geometry_msgs/tf2_geometry_msgs.hpp>


namespace rm_aim
{
  PnPSolver::PnPSolver(
      const std::array<double, 9> &camera_matrix, const std::vector<double> &dist_coeffs)
      : camera_matrix_(cv::Mat(3, 3, CV_64F, const_cast<double *>(camera_matrix.data())).clone()),
        dist_coeffs_(cv::Mat(1, 5, CV_64F, const_cast<double *>(dist_coeffs.data())).clone())
  {
    // Unit: m
    constexpr double small_half_y = SMALL_ARMOR_WIDTH / 2.0 / 1000.0;
    constexpr double small_half_z = SMALL_ARMOR_HEIGHT / 2.0 / 1000.0;
    constexpr double large_half_y = LARGE_ARMOR_WIDTH / 2.0 / 1000.0;
    constexpr double large_half_z = LARGE_ARMOR_HEIGHT / 2.0 / 1000.0;

    // Model coordinate: x forward, y left, z up
    // 左上 → 右上 → 右下 → 左下（与 Armor 的 2D 角点顺序完全一致）
    small_armor_points_.emplace_back(cv::Point3f(0, small_half_y, small_half_z));  // 1. 左上点（y正、z正）
    small_armor_points_.emplace_back(cv::Point3f(0, -small_half_y, small_half_z)); // 2. 右上点（y负、z正）
    small_armor_points_.emplace_back(cv::Point3f(0, -small_half_y, -small_half_z));// 3. 右下点（y负、z负）
    small_armor_points_.emplace_back(cv::Point3f(0, small_half_y, -small_half_z));  // 4. 左下点（y正、z负）

    large_armor_points_.emplace_back(cv::Point3f(0, large_half_y, large_half_z));  // 1. 左上点
    large_armor_points_.emplace_back(cv::Point3f(0, -large_half_y, large_half_z)); // 2. 右上点
    large_armor_points_.emplace_back(cv::Point3f(0, -large_half_y, -large_half_z));// 3. 右下点
    large_armor_points_.emplace_back(cv::Point3f(0, large_half_y, -large_half_z));  // 4. 左下点
  }

  bool PnPSolver::solvePnP(const Armor &armor, cv::Mat &rvec, cv::Mat &tvec)
  {
    std::vector<cv::Point2f> image_armor_points;

    // Fill in image points
    image_armor_points.emplace_back(armor.apex[0]);
    image_armor_points.emplace_back(armor.apex[1]);
    image_armor_points.emplace_back(armor.apex[2]);
    image_armor_points.emplace_back(armor.apex[3]);

    // Solve pnp
    auto object_points = armor.type == ArmorType::SMALL ? small_armor_points_ : large_armor_points_;
    return cv::solvePnP(
        object_points, image_armor_points, camera_matrix_, dist_coeffs_, rvec, tvec, false,
        cv::SOLVEPNP_IPPE);
  }

  float PnPSolver::calculateDistanceToCenter(const cv::Point2f &image_point)
  {
    float cx = camera_matrix_.at<double>(0, 2);
    float cy = camera_matrix_.at<double>(1, 2);
    return cv::norm(image_point - cv::Point2f(cx, cy));
  }

  geometry_msgs::msg::Pose PnPSolver::optimize_yaw_direct(const cv::Mat &rvec, const cv::Mat &tvec, const Armor &armor, const cv::Mat &camera_matrix, const cv::Mat &distort_coeffs) const
  {
    geometry_msgs::msg::Pose optimized_pose;

    // -------------------------- 步骤1：原始rvec→yaw/pitch/roll --------------------------
    cv::Mat rotation_matrix;
    cv::Rodrigues(rvec, rotation_matrix);
    Eigen::Matrix3d R_eigen;
    cv::cv2eigen(rotation_matrix, R_eigen);
    Eigen::Vector3d ypr_raw = R_eigen.eulerAngles(2, 1, 0); // Z-Y-X：yaw/pitch/roll
    double yaw_raw = ypr_raw[0], pitch_raw = ypr_raw[1], roll_raw = ypr_raw[2];

    // -------------------------- 步骤2：搜索最优yaw --------------------------
    constexpr double SEARCH_RANGE = 140;
    double yaw0 = angles::normalize_angle(yaw_raw - SEARCH_RANGE / 2 * CV_PI / 180.0);
    double min_error = 1e10, best_yaw = yaw_raw;

    for (int i = 0; i < SEARCH_RANGE; i++)
    {
      double yaw_candidate = angles::normalize_angle(yaw0 + i * CV_PI / 180.0);
      double error = calculate_reprojection_error_direct(
          yaw_candidate, pitch_raw, roll_raw, tvec, armor, camera_matrix, distort_coeffs);
      if (error < min_error)
      {
        min_error = error;
        best_yaw = yaw_candidate;
      }
    }
    // -------------------------- 步骤3：最优yaw→四元数 --------------------------
    // 最优姿态→旋转矩阵
    Eigen::Matrix3d R_optimized = (Eigen::AngleAxisd(best_yaw, Eigen::Vector3d::UnitZ()) * Eigen::AngleAxisd(pitch_raw, Eigen::Vector3d::UnitY()) * Eigen::AngleAxisd(roll_raw, Eigen::Vector3d::UnitX())).matrix();;
    // 旋转矩阵→tf2四元数
    cv::Mat R_optimized_cv;
    cv::eigen2cv(R_optimized, R_optimized_cv);
    tf2::Matrix3x3 tf2_rot_mat(
        R_optimized_cv.at<double>(0, 0), R_optimized_cv.at<double>(0, 1), R_optimized_cv.at<double>(0, 2),
        R_optimized_cv.at<double>(1, 0), R_optimized_cv.at<double>(1, 1), R_optimized_cv.at<double>(1, 2),
        R_optimized_cv.at<double>(2, 0), R_optimized_cv.at<double>(2, 1), R_optimized_cv.at<double>(2, 2));
    tf2::Quaternion tf2_quat;
    tf2_rot_mat.getRotation(tf2_quat);

    // -------------------------- 步骤4：填充返回的Pose --------------------------
    // 位置：直接复用原始tvec（优化不改变位置）
    optimized_pose.position.x = tvec.at<double>(0);
    optimized_pose.position.y = tvec.at<double>(1);
    optimized_pose.position.z = tvec.at<double>(2);
    // 姿态：优化后的四元数
    optimized_pose.orientation = tf2::toMsg(tf2_quat);

    return optimized_pose;
  }

  double PnPSolver::calculate_reprojection_error_direct(double yaw_candidate, double pitch, double roll, const cv::Mat &tvec, const Armor &armor, const cv::Mat &camera_matrix, const cv::Mat &distort_coeffs) const
  {
    // 1. 构建候选姿态的旋转矩阵
    Eigen::Matrix3d R_candidate = (Eigen::AngleAxisd(yaw_candidate, Eigen::Vector3d::UnitZ()) * 
                                  Eigen::AngleAxisd(pitch, Eigen::Vector3d::UnitY()) * 
                                  Eigen::AngleAxisd(roll, Eigen::Vector3d::UnitX())).matrix();

    // 2. 旋转矩阵 → 候选rvec
    cv::Mat R_candidate_cv, rvec_candidate;
    cv::eigen2cv(R_candidate, R_candidate_cv);
    cv::Rodrigues(R_candidate_cv, rvec_candidate);

    // 3. 装甲板3D模板点（按类型选择）
    const std::vector<cv::Point3f> object_points = (armor.type == ArmorType::LARGE)
                                                       ? large_armor_points_
                                                       : small_armor_points_;

    // 4. 重投影到2D图像
    std::vector<cv::Point2f> reprojected_points;
    cv::projectPoints(object_points, rvec_candidate, tvec, camera_matrix, distort_coeffs, reprojected_points);

    // 5. 计算4个角点的总误差（欧氏距离和）
    double total_error = 0.0;
    for (int i = 0; i < 4; i++)
    {
      total_error += cv::norm(armor.apex[i] - reprojected_points[i]);
    }

    return total_error;
  }

} // namespace rm_aim
